from tqdm import tqdm
import numpy
import shelve
from gene_graph_struct_by_radare2 import radare2_wrapper
from datatype import Tokens, Function, Instruction
import os
import torch
import asm2vec.model
import sys
sys.path.append(
    r"/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master")
sys.path.append(
    r"/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master/asm2vec")

# 有汇编代码的路径
# radara2SamplePath="/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2/"
radara2SamplePath = "/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/radare2/"
# 得到函数嵌入样本的保存路径
# radara2SamplesavePath="/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_predata/"
radara2SamplesavePath = "/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/radare2_predata/"


class gene_embedding():
    """
        使用训练好的asm2vec模型来获得样本的embedding,
        这里符合原论文使用radare处理的汇编指令
    """
    # mpath="/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master/model.pt"
    mpath = "/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master/model2.pt"
    device = "cuda"
    radare2Tools = radare2_wrapper()
    epochs = 30
    lr = 0.02

    def __init__(self) -> None:
        pass

    def cosine_similarity(self, v1, v2):
        return (v1 @ v2 / (v1.norm() * v2.norm())).item()

    def get_embedding(self, discodes):
        """
            输入样本各个函数的反汇编代码
            返回函数对应的embedding,维度为128维
            备注：相同的输入，相似度很高，但是有的时候也会只有80几的相似度
        """
        # 保存基本块和指令
        # 保存的格式中的汇编指令是radare2生成的，作者只加了个title
        # 两个解决思路：
        #   第一条：使用radare2生成样本的图结构，其余的代码不需要变动,方法较难
        #   第二条,使用ida之前处理的图结构和样本,需要进行较多的数据处理,且不太合理,方法可能会简单一点
        #   初步打算使用第一条方法,后续来写
        functions, tokens_new = [], Tokens()
        for discode in discodes:
            # 能正确加载并运行
            fn = Function.load(discode)
            functions.append(fn)
            tokens_new.add(fn.tokens())

        epochs = 30
        lr = 0.02
        # load model, tokens
        model, tokens = asm2vec.utils.load_model(
            self.mpath, device=self.device)
        # functions, tokens_new = asm2vec.utils.load_data(functions)
        tokens.update(tokens_new)
        # 第一个参数是新增加函数的数量
        model.update(len(functions), tokens.size())
        model = model.to(self.device)

        # train function embedding
        model = asm2vec.utils.train(
            functions,
            tokens,
            model=model,
            epochs=epochs,
            device=self.device,
            mode='test',
            learning_rate=lr
        )
        # 这里改成cuda,会快一点嘛
        v1 = model.to(self.device).embeddings_f(
            torch.tensor(list(range(len(functions)))).cuda())
        return v1

    def trans_sample_discode_to_embedding(self, sampleName):
        tar = self.radare2Tools.getter(sampleName)
        name_to_id = tar["name_to_id"]
        discode = tar["discode"]
        size = tar["size"]
        # 转成id_to_name,避免字典修改位置的情况
        id_to_name = {}
        for i in name_to_id.keys():
            id_to_name[name_to_id[i]] = i
        discodes = []
        for i in range(size):
            name = id_to_name[i]
            discodes.append(discode[name])
        att = self.get_embedding(discodes)
        # 这里要移到内存中才能将tensor转成numpy
        tar["att"] = att.cpu().detach().numpy()
        with shelve.open(radara2SamplesavePath+sampleName.split("_")[0]) as file:
            file["data"] = tar
            file.close()


def judge_sample_process_valid(errorName):
    """
        存在12个样本什么都没有,radare2就是不能处理。
        ['742a22390bc72ec38eb18c95446972e8_origin', '2190be4494ab75983c79a1856044e2cf_origin', '974f3efc624ef376d41874f9edd294b2_origin', 'f01afa2e1c0267b2dfb89693a86bb060_origin', '903df3dae25466d58e04caf52df393a9_origin', '78d7fff3916ce40dfe50a50277ff60f8_origin', 'c8e1f8ecb5ba5a0a3c9338c0715abc35_origin', 'bf313223cd0c1bdf938c22eb7344371f_origin', 'b8b4c27058b737bc4149879395bbfa0a_origin', 'ca4812960ec2e217cb8b43f80ba8447b_origin', 'ae0381760bdf4bbcee13c242ce365787_origin', 'f928a7e68d826fce5cc4e87e943341d7_origin']
        确定时radara2不能处理的样本，
        这里来判断一下和之前划分的数据集是否一致
    """
    # 获得划分的样本信息
    sampleLablesPath = "/home/cyw/projects/function_sim_project/all_data/newPair2/sample_and_lables"
    with shelve.open(sampleLablesPath) as file:
        trainName, trainLable = file["trainName"], file["trainLable"]
        testName, testLable = file["testName"], file["testLable"]
        validName, validLable = file["validName"], file["validLable"]
    flag = True
    for name in errorName:
        i = name.split("_")[0]
        if i in trainName or i in testName or i in validName:
            print(i)
            flag = False
    if flag:
        print("未处理的数据不影响模型实验")
    else:
        print("异常数据影响先前模型实验")


if __name__ == "__main__":
    # 获得已经处理好的数据
    # nohup python /home/cyw/projects/function_sim_project/siamese_Graphsage/gene_function_embedding.py
    finishedName = {}
    cnt = 0
    g = os.walk(radara2SamplesavePath)
    for path, dir_list, file_list in g:
        for fileName in tqdm(file_list):
            name, hz = fileName.split(".")
            if hz == "dir":
                finishedName[name] = True
                cnt += 1
    print("已经生成了{}个样本的函数嵌入\n".format(cnt))
    a = gene_embedding()
    error = 0
    errorName = []
    g = os.walk(radara2SamplePath)
    for path, dir_list, file_list in g:
        for fileName in tqdm(file_list):
            name, hz = fileName.split(".")
            if hz == "dir" and len(name.split("_")) == 2 and name.split("_")[0] not in finishedName:
                try:
                    a.trans_sample_discode_to_embedding(name)
                except Exception as e:
                    print(e)
                    error += 1
                    errorName.append(name)
    print("生成函数嵌入完成！!")
    print("共{}个样本未能处理".format(error))
    print(errorName)
    judge_sample_process_valid(errorName)
